import dgl
import json
import itertools
import numpy as np
from tqdm import tqdm

import torch
from torch.utils.data import Dataset, DataLoader


class tsp_dataset(Dataset):
    def __init__(self, data_path):
        self.graph_list = []
        self.edge_labels = []
        # 读取json文件
        with open(data_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        #data可以包含多组数据，一组数据其实就是一组图数据，包含四个部分：节点坐标，任意两点之间的距离，边的起点和终点，边的标签是否选择这条路线
        for i in tqdm(range(len(data))):
            # 数据预处理
            i = str(i)
            # data数据第i个图的ndata 信息， 12个node
            num_nodes = len(data[i]['ndata'])
            # edges=12 * 11 = 132
            edges = torch.tensor(data[i]['edges'], dtype=torch.int64)
            # n_feat 12 *2 tensor
            n_feat = torch.tensor(data[i]['ndata'], dtype=torch.float32)
            # e_feat 132 *1 tensor
            e_feat = torch.tensor(data[i]['edata'], dtype=torch.float32).view(-1, 1)
            # edge_label list 132
            edge_label = data[i]['label']
            # 构造图
            graph = dgl.DGLGraph()
            graph.add_nodes(num_nodes)
            # add_edges 需要传入两个list，表示边的起点和终点 edges[:, 0], edges[:, 1] 分别表示边的起点和终点
            graph.add_edges(edges[:, 0], edges[:, 1])
            # 节点特征 n_feat 12*2  边特征 e_feat 132*1
            graph.ndata['feat'] = n_feat
            graph.edata['feat'] = e_feat
            # 保存数据
            self.graph_list.append(graph)
            # edge_labels 边分类标签
            self.edge_labels.append(edge_label)

    def __getitem__(self, idx):
        return self.graph_list[idx], self.edge_labels[idx]

    def __len__(self):
        return len(self.graph_list)


class collate_fn(object):
    def __init__(self):
        super(collate_fn, self).__init__()

    def __call__(self, batch_data):
        # The input samples is a list of pairs (graph, label).
        # 输入是一个mini_batch的list, (graph, label)
        graphs, labels = map(list, zip(*batch_data))
        # Edge classification labels need to be flattened to 1D lists
        # 真实标签需要进行展平
        labels = torch.LongTensor(np.array(list(itertools.chain(*labels))))

        batched_graph = dgl.batch(graphs)

        return batched_graph, labels


if __name__ == '__main__':
    data_path = "../dataset/small/val_data.json"
    data_path = "../dataset/small/train_data.json"
    data_set = tsp_dataset(data_path)
    data_loader = DataLoader(data_set, batch_size=16, shuffle=True, collate_fn=collate_fn())

    for data_iter in data_loader:
        batch_graph, edge_labels = data_iter
        print(batch_graph)
